"""
连接管理器模块
使用contextmanager的方式管理SSH隧道和数据库连接
增强版本：添加重试机制、连接预检和错误恢复，包含StarRocks连接补丁
"""
# 🔧 重要：应用StarRocks连接补丁（必须在其他数据库相关导入之前）
try:
    from airflow_connection_patch import apply_all_patches
    patch_success = apply_all_patches()
    print(f"🎉 StarRocks连接补丁应用结果: {patch_success}")
except Exception as e:
    print(f"❌ 导入StarRocks连接补丁失败: {str(e)}")
    # 继续执行，但记录警告
    print("⚠️ 将继续使用原始连接方式，可能会遇到超时问题")

import pymysql
import traceback
import os
import time
import random
from contextlib import contextmanager
from sshtunnel import SSHTunnelForwarder
from airflow.hooks.base import BaseHook
from airflow.models import Variable
import logging

logger = logging.getLogger(__name__)

# 数据目录配置 - 集中管理所有路径配置
current_dir = os.path.dirname(os.path.abspath(__file__))
BASE_DATA_DIR = os.path.join(current_dir, "data")
FAQ_DATA_DIR = os.path.join(BASE_DATA_DIR, "faq")
EC_DATA_DIR = os.path.join(BASE_DATA_DIR, "errorcode")

# Dify配置
local_dify = Variable.get("local_dify")
aws_dify = Variable.get("aws_dify")
agi_email_password = Variable.get("agi_email_password")
DIFY_KNOWLEDGE_DATA = {
    'local': local_dify,
    'aws': aws_dify
}

EMAIL_CONFIG = {
    'imap_server': 'imap.poweroak.net',
    'imap_port': 993,
    'smtp_server': 'smtp.poweroak.net',
    'smtp_port': 994,
    'username': 'baihao@poweroak.net',
    'password': agi_email_password,
    'sender': '请勿回复 <bluetti-agi@poweroak.net>',
    'receivers': [
        'guote@poweroak.net',
        'hujin@poweroak.net',
        'shiyuxin@poweroak.net',
        ]  # Add all receiver email addresses here
}

logger.info(f"📁 数据目录配置:")
logger.info(f"  📂 基础数据目录: {BASE_DATA_DIR}")
logger.info(f"  📂 FAQ数据目录: {FAQ_DATA_DIR}")
logger.info(f"  📂 错误代码数据目录: {EC_DATA_DIR}")

def get_ssh_connection():
    """从Airflow connection获取SSH隧道参数"""
    logger.info("🔗 开始获取SSH连接配置...")
    try:
        ssh_conn = BaseHook.get_connection('product-cloud-app')
        config = {
            'host': ssh_conn.host,
            'port': ssh_conn.port or 22,
            'username': ssh_conn.login,
            'password': ssh_conn.password
        }
        logger.info(f"✅ SSH连接配置获取成功: {config['username']}@{config['host']}:{config['port']}")
        return config
    except Exception as e:
        logger.error(f"❌ 获取SSH连接配置失败: {str(e)}")
        logger.error(f"🔍 详细错误信息:\n{traceback.format_exc()}")
        raise

def get_database_connection(connection_id='aliyun-database'):
    """从Airflow connection获取数据库连接参数"""
    logger.info(f"🗄️ 开始获取数据库连接配置: {connection_id}")
    try:
        db_conn = BaseHook.get_connection(connection_id)
        config = {
            'host': db_conn.host,
            'port': db_conn.port or 3306,
            'user': db_conn.login,
            'password': db_conn.password,
            'database': db_conn.schema
        }
        logger.info(f"✅ 数据库连接配置获取成功: {config['user']}@{config['host']}:{config['port']}/{config['database']}")
        return config
    except Exception as e:
        logger.error(f"❌ 获取数据库连接配置失败 ({connection_id}): {str(e)}")
        logger.error(f"🔍 详细错误信息:\n{traceback.format_exc()}")
        raise

def test_ssh_connectivity(ssh_config, max_retries=3):
    """测试SSH连接可用性"""
    for attempt in range(max_retries):
        try:
            logger.info(f"🔧 测试SSH连接 (尝试 {attempt + 1}/{max_retries})...")
            tunnel = SSHTunnelForwarder(
                (ssh_config['host'], ssh_config['port']),
                ssh_username=ssh_config['username'],
                ssh_password=ssh_config['password'],
                remote_bind_address=('127.0.0.1', 22),  # 测试连接
                local_bind_address=('127.0.0.1', 0)  # 自动分配端口
            )
            tunnel.start()
            logger.info(f"✅ SSH连接测试成功 (尝试 {attempt + 1})")
            tunnel.close()
            return True
        except Exception as e:
            logger.warning(f"⚠️ SSH连接测试失败 (尝试 {attempt + 1}): {str(e)}")
            if attempt < max_retries - 1:
                wait_time = (attempt + 1) * 2 + random.uniform(0, 2)
                logger.info(f"⏳ 等待 {wait_time:.1f} 秒后重试...")
                time.sleep(wait_time)
            else:
                logger.error(f"❌ SSH连接测试失败，已达到最大重试次数")
                return False
    return False

@contextmanager
def get_database_cursor(db_connection_id='aliyun-database', max_retries=3, enable_precheck=True):
    """
    Context manager for database connection with SSH tunnel
    增强版本：添加重试机制、连接预检和错误恢复
    
    Args:
        db_connection_id: 数据库连接ID
            - 'aliyun-database-faq': FAQ爬虫使用
            - 'aliyun-database-ec-ts': 故障排除爬虫使用
        max_retries: 最大重试次数
        enable_precheck: 是否启用SSH连接预检
    """
    logger.info(f"🚀 开始建立数据库连接 (使用连接: {db_connection_id}, 重试次数: {max_retries})...")
    
    ssh_config = get_ssh_connection()
    db_config = get_database_connection(db_connection_id)
    
    # SSH连接预检
    # if enable_precheck:
    #     logger.info("🔍 执行SSH连接预检...")
    #     if not test_ssh_connectivity(ssh_config):
    #         raise Exception("SSH连接预检失败，无法继续")
    #     logger.info("✅ SSH连接预检通过")
    
    tunnel = None
    connection = None
    last_exception = None
    
    for attempt in range(max_retries):
        try:
            logger.info(f"🔧 建立SSH隧道连接 (尝试 {attempt + 1}/{max_retries})...")
            
            # 清理之前的连接
            if tunnel:
                try:
                    tunnel.close()
                except:
                    pass
            if connection:
                try:
                    connection.close()
                except:
                    pass
            
            # 建立SSH隧道连接
            tunnel = SSHTunnelForwarder(
                (ssh_config['host'], ssh_config['port']),
                ssh_username=ssh_config['username'],
                ssh_password=ssh_config['password'],
                remote_bind_address=(db_config['host'], db_config['port']),
                # local_bind_address=('127.0.0.1', 0)  # 自动分配端口避免冲突
            )
            
            tunnel.start()
            logger.info(f"✅ SSH隧道已建立 - 本地端口: {tunnel.local_bind_port} (尝试 {attempt + 1})")
            logger.info(f"🔌 正在连接数据库 ({db_connection_id}) (尝试 {attempt + 1})...")
            # 使用隧道连接数据库
            connection = pymysql.connect(
                host='127.0.0.1',
                port=tunnel.local_bind_port,
                user=db_config['user'],
                password=db_config['password'],
                database=db_config['database'],
                # charset='utf8mb4',
                connect_timeout=15,  # 增加超时时间
                # read_timeout=30,
                # write_timeout=30,
                # autocommit=True,  # 启用自动提交
                # init_command="SET SESSION wait_timeout=3600"  # 设置会话超时
            )
            
            logger.info(f"✅ 数据库连接成功 ({db_connection_id}) (尝试 {attempt + 1})")
            
            # 返回游标
            with connection.cursor() as cursor:
                # logger.info(f"📊 数据库游标已创建，准备执行操作 ({db_connection_id})...")
                yield cursor
                # logger.info(f"✅ 数据库操作完成 ({db_connection_id})")
                return  # 成功完成，退出重试循环
                
        except Exception as e:
            last_exception = e
            logger.error(f"❌ 数据库连接出现错误 ({db_connection_id}) (尝试 {attempt + 1}): {str(e)}")
            
            if attempt < max_retries - 1:
                wait_time = (attempt + 1) * 3 + random.uniform(1, 3)
                logger.info(f"⏳ 等待 {wait_time:.1f} 秒后重试...")
                time.sleep(wait_time)
            else:
                logger.error(f"❌ 数据库连接失败，已达到最大重试次数 ({max_retries})")
                logger.error(f"🔍 最后一次错误详情:\n{traceback.format_exc()}")
        finally:
            # 每次尝试后都清理连接，避免连接泄露
            if connection and attempt < max_retries - 1:  # 不是最后一次尝试
                try:
                    connection.close()
                    logger.info(f"🔐 清理数据库连接 (尝试 {attempt + 1})")
                except Exception as e:
                    logger.warning(f"⚠️ 清理数据库连接时出现警告: {str(e)}")
                    
            if tunnel and attempt < max_retries - 1:  # 不是最后一次尝试
                try:
                    tunnel.close()
                    logger.info(f"🔒 清理SSH隧道 (尝试 {attempt + 1})")
                except Exception as e:
                    logger.warning(f"⚠️ 清理SSH隧道时出现警告: {str(e)}")
    
    # 所有重试都失败了，最终清理并抛出异常
    if connection:
        try:
            connection.close()
            logger.info(f"🔐 最终清理：数据库连接已关闭 ({db_connection_id})")
        except Exception as e:
            logger.warning(f"⚠️ 最终清理数据库连接时出现警告 ({db_connection_id}): {str(e)}")
            
    if tunnel:
        try:
            tunnel.close()
            logger.info("🔒 最终清理：SSH隧道已关闭")
        except Exception as e:
            logger.warning(f"⚠️ 最终清理SSH隧道时出现警告: {str(e)}")
    
    # 抛出最后一个异常
    if last_exception:
        raise last_exception
    else:
        raise Exception(f"数据库连接失败，原因未知 ({db_connection_id})")

# 为了向后兼容，保留原函数名
def get_faq_database_cursor():
    """FAQ爬虫专用数据库连接"""
    return get_database_cursor('aliyun-database-faq')

def get_ts_ec_database_cursor():
    """故障排除爬虫专用数据库连接"""
    return get_database_cursor('aliyun-database-ec-ts')
